#include "inject.h"
#include <tlhelp32.h>
#include <psapi.h>
#include <shlwapi.h>
#include "overlay/overlay.h"

bool CheckProcess(DWORD pID)
{
  bool hasModule = false;
  HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pID);
  if (hSnapshot != INVALID_HANDLE_VALUE)
  {
    MODULEENTRY32 me;
    me.dwSize = sizeof me;
    if (Module32First(hSnapshot, &me))
    {
      do
      {
        if (!wcsicmp(me.szModule, L"d3d9.dll") ||
            !wcsicmp(me.szModule, L"d3d10.dll") ||
            !wcsicmp(me.szModule, L"d3d11.dll"))
          hasModule = true;
      } while (Module32Next(hSnapshot, &me));
    }
    CloseHandle(hSnapshot);
  }
  return hasModule;
}
static HICON GetWindowIcon(HWND hWnd)
{
  DWORD result;
  if (SendMessageTimeout(hWnd, WM_GETICON, ICON_SMALL, 0, SMTO_NORMAL | SMTO_ABORTIFHUNG, 250, &result) && result)
    return (HICON) result;
  if (SendMessageTimeout(hWnd, WM_GETICON, ICON_BIG, 0, SMTO_NORMAL | SMTO_ABORTIFHUNG, 250, &result) && result)
    return (HICON) result;
  if (result = GetClassLong(hWnd, GCL_HICONSM))
    return (HICON) result;
  return (HICON) GetClassLong(hWnd, GCL_HICON);
}
void GetProcessInfo(HWND hWnd, ProcessEntry& pe)
{
  pe.is64 = FALSE;
  pe.tID = GetWindowThreadProcessId(hWnd, &pe.pID);
  GetWindowText(hWnd, pe.name, MAX_PATH);
  if (pe.name[0] == 0)
  {
    HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
      FALSE, pe.pID);
    if (hProcess)
    {
      if (GetModuleFileNameEx(hProcess, NULL, pe.name, MAX_PATH))
        PathStripPath(pe.name);
      CloseHandle(hProcess);
    }
  }
  pe.hIcon = GetWindowIcon(hWnd);
}
static BOOL CALLBACK EnumWindowsCallback(HWND hWnd, LPARAM lParam)
{
  if (!IsWindowVisible(hWnd)) return TRUE;
  Array<ProcessEntry>* list = (Array<ProcessEntry>*) lParam;
  DWORD pID;
  DWORD tID = GetWindowThreadProcessId(hWnd, &pID);
  for (int i = 0; i < list->length(); i++)
    if ((*list)[i].pID == pID)
      return TRUE;
  if (CheckProcess(pID))
  {
    ProcessEntry& p = list->push();
    GetProcessInfo(hWnd, p);
  }
  return TRUE;
}

void EnumProcesses(Array<ProcessEntry>& processes)
{
  EnumWindows(EnumWindowsCallback, (LPARAM) &processes);
}

DWORD GetProcessThread(DWORD pID)
{
  DWORD tID = 0;
  HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
  if (hSnapshot != INVALID_HANDLE_VALUE)
  {
    THREADENTRY32 te;
    te.dwSize = sizeof te;
    if (Thread32First(hSnapshot, &te))
    {
      do
      {
        if (tID == 0 && te.th32OwnerProcessID == pID)
          tID = te.th32ThreadID;
      } while (Thread32Next(hSnapshot, &te));
    }
    CloseHandle(hSnapshot);
  }
  return tID;
}

HHOOK InstallHook(DWORD tID)
{
  HMODULE hDLL = GetModuleHandle(L"kappa32");
  if (tID && hDLL)
    return SetWindowsHookEx(WH_CBT, HookProc, hDLL, tID);
  return NULL;
}
HHOOK InstallMouseHook(DWORD tID)
{
  HMODULE hDLL = GetModuleHandle(L"kappa32");
  if (tID && hDLL)
    return SetWindowsHookEx(WH_MOUSE, MouseHookProc, hDLL, tID);
  return NULL;
}
HHOOK InstallKeyboardHook(DWORD tID)
{
  HMODULE hDLL = GetModuleHandle(L"kappa32");
  if (tID && hDLL)
    return SetWindowsHookEx(WH_KEYBOARD, KeyboardHookProc, hDLL, tID);
  return NULL;
}
void RemoveHook(HHOOK hHook)
{
  if (hHook)
    UnhookWindowsHookEx(hHook);
}
